package net.cyue.ort.llm.template;

import com.hubspot.jinjava.Jinjava;
import com.hubspot.jinjava.JinjavaConfig;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.TimeUnit;

public class ModelChatTemplate {

    private final String template;
    private final String bosToken;
    private final String eosToken;
    private final boolean debug;
    
    // 静态单例Jinjava引擎
    private static volatile Jinjava ENGINE = null;
    private static final Object LOCK = new Object();
    
    // 模板渲染结果缓存
    private static final Cache<String, String> templateCache = CacheBuilder.newBuilder()
            .maximumSize(1000)
            .expireAfterWrite(10, TimeUnit.MINUTES)
            .build();

    public ModelChatTemplate(
            String template,
            String bosToken,
            String eosToken
    ) {
        this(
            template,
            bosToken,
            eosToken,
            false
        );
    }

    public ModelChatTemplate(
        String template,
        String bosToken,
        String eosToken,
        boolean debug
    ) {
        this.template = template;
        this.bosToken = bosToken;
        this.eosToken = eosToken;
        this.debug = debug;
        initializeEngine();
    }

    private static void initializeEngine() {
        if (ENGINE == null) {
            synchronized (LOCK) {
                if (ENGINE == null) {
                    ENGINE = new Jinjava(
                        JinjavaConfig.newBuilder()
                            .withTrimBlocks(true)
                            .withLstripBlocks(true)
                            .withNestedInterpretationEnabled(true)
                            .withFailOnUnknownTokens(false)
                            .build()
                    );
                }
            }
        }
    }

    public String apply(List<ModelChatMessage> messages) {
        return apply(messages, true);
    }

    public String apply(String[] messages) {
        return apply(messages, true);
    }

    public String apply(List<ModelChatMessage> messages, boolean addGenerationPrompt) {
        try {
            // 生成缓存key
            String cacheKey = generateCacheKey(messages, addGenerationPrompt);

            // 尝试从缓存获取
            String cachedResult = templateCache.getIfPresent(cacheKey);
            if (cachedResult != null) {
                return cachedResult;
            }

            Map<String, Object> context = new HashMap<>();
            context.put("messages", convertMessages(messages));
            context.put("bos_token", bosToken);
            context.put("eos_token", eosToken);
            context.put("add_generation_prompt", addGenerationPrompt);

            String result = renderTemplate(context);
            
            // 存入缓存
            templateCache.put(cacheKey, result);
            
            return result;
        } catch (Exception e) {
            throw new TemplateRenderException("Failed to render template", e);
        }
    }

    public String apply(String[] messages, boolean addGenerationPrompt) {
        try {
            String cacheKey = generateCacheKey(messages, addGenerationPrompt);
            
            String cachedResult = templateCache.getIfPresent(cacheKey);
            if (cachedResult != null) {
                return cachedResult;
            }

            Map<String, Object> context = new HashMap<>();
            context.put("messages", convertMessages(messages));
            context.put("bos_token", bosToken);
            context.put("eos_token", eosToken);
            context.put("add_generation_prompt", addGenerationPrompt);

            String result = renderTemplate(context);
            templateCache.put(cacheKey, result);
            
            return result;
        } catch (Exception e) {
            throw new TemplateRenderException("Failed to render template", e);
        }
    }

    private String renderTemplate(Map<String, Object> context) {
        try {
            String result = ENGINE.render(template, context).trim();
            
            if (this.debug) {
                System.out.println("\ntemplate: " + template);
                System.out.println("result: " + result);
            }
            return result;
        } catch (Exception e) {
            throw new TemplateRenderException("Template rendering failed", e);
        }
    }

    private String generateCacheKey(Object messages, boolean addGenerationPrompt) {
        return String.format("%s_%s_%b", 
            messages.hashCode(), 
            template.hashCode(), 
            addGenerationPrompt
        );
    }

    private List<Map<String, String>> convertMessages(List<ModelChatMessage> messages) {
        List<Map<String, String>> result = new ArrayList<>();
        for (ModelChatMessage m : messages) {
            Map<String, String> messageMap = new HashMap<>();
            messageMap.put("role", m.role.toString().toLowerCase());
            messageMap.put("content", m.content);
            result.add(messageMap);
        }
        return result;
    }

    public List<Map<String, String>> convertMessages(String[] messages) {
        List<Map<String, String>> result = new ArrayList<>();
        for (String m : messages) {
            Map<String, String> messageMap = new HashMap<>();
            messageMap.put("content", m);
            result.add(messageMap);
        }
        return result;
    }

    public static class TemplateRenderException extends RuntimeException {
        public TemplateRenderException(String message) {
            super(message);
        }

        public TemplateRenderException(String message, Throwable cause) {
            super(message, cause);
        }
    }

    // 构建器模式
    public static class Builder {
        private String template;
        private String bosToken = "";
        private String eosToken = "";
        private boolean debug = false;

        public Builder template(String template) {
            this.template = template;
            return this;
        }

        public Builder bosToken(String bosToken) {
            this.bosToken = bosToken;
            return this;
        }

        public Builder eosToken(String eosToken) {
            this.eosToken = eosToken;
            return this;
        }

        public Builder debug(boolean debug) {
            this.debug = debug;
            return this;
        }

        public ModelChatTemplate build() {
            if (template == null || template.isEmpty()) {
                throw new IllegalStateException("Template cannot be null or empty");
            }
            return new ModelChatTemplate(template, bosToken, eosToken, debug);
        }
    }
}


